from PySide6.QtCore import Qt
from PySide6.QtWidgets import (QFrame, QVBoxLayout, QHBoxLayout, 
                              QFileDialog, QMessageBox, QTableWidget,
                              QTableWidgetItem, QLabel, QHeaderView,
                              QAbstractItemView)
from qfluentwidgets import (PrimaryPushButton, LineEdit, InfoBar, 
                           InfoBarPosition)
import pandas as pd
from utils.template_engine_v2 import TemplateEngine
import gc

class ReportGenerator(QFrame):
    """ Excel导入生成Word报告界面 """
    
    def __init__(self, title: str, parent=None):
        super().__init__(parent=parent)
        self.title = title
        self.excel_path = None
        self.template_engine = TemplateEngine()
        
        # 创建布局
        self.vBoxLayout = QVBoxLayout(self)
        self.hBoxLayout = QHBoxLayout()
        
        # 创建控件
        self.importButton = PrimaryPushButton('导入Excel', self)
        self.pathEdit = LineEdit(self)
        self.generateButton = PrimaryPushButton('生成报告', self)
        
        # 创建预览表格
        self.previewTable = QTableWidget(self)
        self.previewTable.setColumnCount(2)
        self.previewTable.setHorizontalHeaderLabels(['自定义变量', '值'])
        
        # 添加模板文件选择
        self.templateLabel = QLabel('Word模板:', self)
        self.templatePathEdit = LineEdit(self)
        self.templatePathEdit.setPlaceholderText('Word模板文件路径...')
        self.templatePathEdit.setReadOnly(True)
        self.selectTemplateButton = PrimaryPushButton('选择模板', self)
        
        # 初始化UI
        self.initUI()
        self.initSignals()

    def closeEvent(self, event):
        """窗口关闭事件"""
        try:
            # 清理模板引擎资源
            if hasattr(self, 'template_engine'):
                self.template_engine.cleanup()
            
            # 清理预览表格
            if hasattr(self, 'previewTable'):
                self.previewTable.clear()
                self.previewTable.setRowCount(0)
                self.previewTable.setColumnCount(0)
            
            # 强制进行垃圾回收
            gc.collect()
            
            # 接受关闭事件
            event.accept()
            
        except Exception as e:
            print(f"关闭窗口时出错: {str(e)}")
            event.accept()

    def __del__(self):
        """析构函数"""
        try:
            # 确保资源被清理
            if hasattr(self, 'template_engine'):
                self.template_engine.cleanup()
            gc.collect()
        except:
            pass

    def initUI(self):
        """ 初始化界面 """
        # 设置布局
        # self.vBoxLayout.setContentsMargins(20, 0, 20, 0)
        self.vBoxLayout.setSpacing(20)
        
        # 设置路径输入框
        self.pathEdit.setPlaceholderText('Excel文件路径...')
        self.pathEdit.setReadOnly(True)
        
        # 设置预览表格
        self.previewTable.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        self.previewTable.setAlternatingRowColors(True)
        self.previewTable.setEditTriggers(QAbstractItemView.NoEditTriggers)
        
        # 添加模板文件选择区域
        templateFileLayout = QHBoxLayout()
        templateFileLayout.addWidget(self.templateLabel)
        templateFileLayout.addWidget(self.templatePathEdit)
        templateFileLayout.addWidget(self.selectTemplateButton)
        
        # 文件选择区域
        fileLayout = QHBoxLayout()
        fileLayout.addWidget(self.pathEdit)
        fileLayout.addWidget(self.importButton)
        
        # 更新布局顺序
        self.vBoxLayout.addLayout(templateFileLayout)
        self.vBoxLayout.addLayout(fileLayout)
        self.vBoxLayout.addWidget(self.previewTable)
        self.vBoxLayout.addWidget(self.generateButton, 0, Qt.AlignCenter)
        
        # 设置对象名称
        self.setObjectName(self.title.replace(' ', '-'))

    def initSignals(self):
        """ 初始化信号连接 """
        self.importButton.clicked.connect(self.importExcel)
        self.generateButton.clicked.connect(self.generateReport)
        self.selectTemplateButton.clicked.connect(self.selectTemplate)

    def updatePreviewTable(self, variables, df=None):
        """更新预览表格"""
        self.previewTable.setRowCount(len(variables))
        
        for i, var in enumerate(variables):
            
            # 创建标签项
            label_item = QTableWidgetItem(f"{var}")
            label_item.setFlags(label_item.flags() & ~Qt.ItemIsEditable)
            self.previewTable.setItem(i, 0, label_item)
            
            # 计算结果
            result = "请输入自定义值..."
            
            # 创建值项
            value_item = QTableWidgetItem(result)
            self.previewTable.setItem(i, 1, value_item)
        # 允许编辑第二列
        self.previewTable.setEditTriggers(QAbstractItemView.DoubleClicked | QAbstractItemView.SelectedClicked)

    def selectTemplate(self):
        """ 选择Word模板文件 """
        file_path, _ = QFileDialog.getOpenFileName(
            self,
            "选择Word模板",
            "",
            "Word Files (*.docx)"
        )
        
        if file_path:
            success, result = self.template_engine.load_template(file_path)
            if success:
                self.templatePathEdit.setText(file_path)
                custom_placeholders = result.get('custom_', set())
                self.updatePreviewTable(custom_placeholders)
            else:
                QMessageBox.warning(self, '错误', result[0])

    def importExcel(self):
        """导入Excel文件"""
        file_path, _ = QFileDialog.getOpenFileName(
            self,
            "选择Excel文件",
            "",
            "Excel Files (*.xlsx *.xls)"
        )
        
        if file_path:
            try:
                self.excel_path = file_path
                self.pathEdit.setText(file_path)
                
                # 读取Excel并更新预览
                self.template_engine.load_excel(file_path)               
                
                # self.updatePreviewTable( df)
                
                InfoBar.success(
                    title='成功',
                    content='Excel文件导入成功',
                    parent=self,
                    position=InfoBarPosition.TOP,
                    duration=2000
                )
                
            except Exception as e:
                QMessageBox.warning(self, '错误', f'导入Excel失败：{str(e)}')

    def getCustomVariableValues(self):
        """从预览表格中获取用户输入的自定义变量值"""
        custom_vars = {}
        for row in range(self.previewTable.rowCount()):
            var_name = self.previewTable.item(row, 0).text()
            var_value = self.previewTable.item(row, 1).text() if self.previewTable.item(row, 1) else ""
            custom_vars[var_name] = var_value
        return custom_vars
    
    def generateReport(self):
        """生成报告"""
        if not self.templatePathEdit.text() or not self.pathEdit.text():
            InfoBar.warning(
                title='警告',
                content='请先选择模板和Excel文件',
                parent=self,
                position=InfoBarPosition.TOP,
                duration=2000
            )
            return
        
        try:
            save_path, _ = QFileDialog.getSaveFileName(
                self,
                "保存报告",
                "",
                "Word Files (*.docx)"
            )

            # 获取用户输入的自定义变量值
            custom_vars = self.getCustomVariableValues()

            # 定义一个函数来提供自定义变量的值
            def custom_var_extractor(custom_vars_set):
                return {var: custom_vars.get(var, "") for var in custom_vars_set}

            # 设置 TemplateEngine 的自定义变量提取器
            self.template_engine.custom_var_extractor = custom_var_extractor
            
            if save_path:
                success, message = self.template_engine.save_document(save_path)
                if success:
                    InfoBar.success(
                        title='成功',
                        content='报告生成成功',
                        parent=self,
                        position=InfoBarPosition.TOP,
                        duration=2000
                    )
                else:
                    raise Exception(message)
                    
        except Exception as e:
            QMessageBox.warning(self, '错误', f'生成报告失败：{str(e)}')

    def cleanup(self):
        """清理资源但保留路径"""
        try:
            if hasattr(self, 'template_engine'):
                # 保存当前路径
                template_path = self.template_engine.template_path
                excel_path = self.pathEdit.text()
                
                # 清理资源
                self.template_engine.cleanup()
                
                # 恢复路径
                self.template_engine.template_path = template_path
                if excel_path:
                    self.pathEdit.setText(excel_path)
                    
        except Exception as e:
            print(f"清理资源时出错: {str(e)}") 